#!/usr/bin/env python3
# @Date    : 2023-03-13
# @Author  : Bright (brt2@qq.com)
# @Link    : https://gitee.com/brt2
# @Version : v0.1.5

from io import BytesIO
import base64

from PIL import Image
import numpy as np

def bytes2pimg(io_bytes):
    if isinstance(io_bytes, bytes):
        io_bytes = BytesIO(io_bytes)

    pimg = Image.open(io_bytes)
    return pimg

def bytes2array(io_bytes):
    pimg = bytes2pimg(io_bytes)
    im_arr = pimg2array(pimg)
    return im_arr

def pimg2bytes(pimg):
    _io = BytesIO()
    pimg.save(_io, "WEBP")
    return _io.getvalue()

def array2pimg(im_arr, mode=None):
    """ wrapper to: PIL.Image.fromarray(im_arr, mode=None)
        这里设定mode为显式参数，因为无法通过channel完全确定mode：
        * 2dim: "1", "L", "P", "I", "F"
        * 3dim: "RGB", "BGR"
        * 4dim: "RGBA", "CMYK", "YCbCr"
    """
    if mode == "BGR":
        NotImplementedError("尚未支持OpenCV格式的图像")
    elif mode is None:
        if im_arr.ndim < 3:
            mode = "L"
        elif im_arr.shape[2] == 3:
            mode = "RGB"  # 无法区分BGR (OpenCV)
        elif im_arr.shape[2] == 4:
            mode = "RGBA"
        else:
            raise Exception("未知的图像类型")
    return Image.fromarray(im_arr, mode)

def pimg2array(pimg):
    """ wrapper to: np.asarray(pil_img) """
    return np.asarray(pimg)

def array2bytes(im_arr):
    # Don't im_arr.tobytes()
    pimg = array2pimg(im_arr)
    return pimg2bytes(pimg)

def array2block(im_arr):
    """ 对ndarray数组直接序列化为bytes结构体
        return: b"{shape_h}_{shape_w}_bytesImage"
    """
    bytes_img = im_arr.tobytes()
    h, w = im_arr.shape
    bytes_data = f"{h}_{w}_".encode() + bytes_img
    return bytes_data

def block2array(data, dtype="uint8"):
    h, w, bytes_img = data.split(b"_", maxsplit=2)
    # h, w = [int(x) for x in size.split(b"x")]
    im_shape = [int(x) for x in [h, w]]
    im_depth = np.frombuffer(bytes_img, dtype=dtype).reshape(im_shape)
    return im_depth

def depth2bytes__(im_depth):
    """ 深度图(dtype=uint16)转bytes流 """
    # 方式1:
    # bytes_img = self.wrap_img_data(im_depth)  # 90Kib, 但由于format转换，dtype变为uint8，数据错乱

    # 方式2，指定16位图片格式
    # pimg = Image.fromarray(im_depth, "I;16")
    # print(">>>", pimg.mode, im_depth.dtype)
    # bytes_img = cvs.pimg2bytes(pimg)

    # 方式3，原始bytes，可以用，但低效:
    bytes_img = array2block(im_depth)

    # 方式4:
    # success, im_encoded = cv2.imencode(".jpg", im_depth)
    # bytes_img = im_encoded.tobytes()  # 20 Kib，但是黑白的二值图，错误

    # 方式5:
    # 可视化深度图：(640, 360), uint16 -> (3, 640, 360), uint8
    # depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(im_depth, alpha=0.03), cv2.COLORMAP_JET)

    # 方式6，PNG压缩图
    # 使用 "I;16" 格式压缩图像
    # bytes_img = self.device.rpc_depth_image()
    return bytes_img

def bytes2base64(bytes_img):
    return base64.b64encode(bytes_img).decode()

base64_frombytes = bytes2base64

def base64_topil(str_img):
    bytes_img = base64.b64decode(str_img)
    return bytes2pimg(bytes_img)
